package com.gogant.spider.http;

import com.gogant.spider.core.Response;
import com.gogant.spider.util.FileUtils;
import java.awt.Image;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.net.HttpURLConnection;
import java.nio.charset.Charset;
import java.util.List;
import javax.imageio.ImageIO;

public class HttpResponse extends Response {
    private int code;
    private String message;
    private byte[] contentBytes;
    private String charset;

    public HttpResponse(HttpRequest req, HttpURLConnection conn) throws Exception {
        super(req, conn);
        code = conn.getResponseCode();
        if (code != -1) {
            message = conn.getResponseMessage();
        }
    }

    public int getCode() {
        return code;
    }

    public boolean isOK() {
        return code == HttpURLConnection.HTTP_OK;
    }

    public String getMessage() {
        return message;
    }

    public byte[] getContentBytes() {
        return contentBytes;
    }

    protected void setContentBytes(byte[] bytes) {
        this.contentBytes = bytes;
    }

    public String getCharset() {
        return charset;
    }

    public void setCharset(String charset) {
        this.charset = charset;
    }

    public InputStream getStreamContent() {
        if (contentBytes == null) {
            return null;
        }
        return new ByteArrayInputStream(contentBytes);
    }

    public String getTextContent() {
        Object content = getContent();
        if (content != null) {
            if (content instanceof String) {
                return (String) content;
            } else {
                return null;
            }
        }
        if (charset == null) {
            charset = Charset.defaultCharset().name();
        }
        try {
            String text = new String(contentBytes, charset);
            setContent(text);
            return text;
        } catch (UnsupportedEncodingException e) {
            return null;
        }
    }

    public Image getImageContent() {
        Object content = getContent();
        if (content != null) {
            if (content instanceof Image) {
                return (Image) content;
            } else {
                return null;
            }
        }

        try {
            Image image = ImageIO.read(new ByteArrayInputStream(contentBytes));
            setContent(image);
            return image;
        } catch (Exception e) {
            return null;
        }
    }

    public <T> T getBean(Class<T> clazz) throws Exception {
        return HtmlBeanHelper.getBean(clazz, this);
    }

    public <T> List<T> getBeanList(Class<T> clazz) throws Exception {
        return HtmlBeanHelper.getBeanList(clazz, this);
    }

    public boolean select(HtmlSelector... selectors) {
        return HtmlHelper.parseHtml(getTextContent(), selectors);
    }

    public boolean saveToFile(File file) {
        return FileUtils.saveToFile(this, file);
    }

    public boolean saveToPath(File path) {
        return FileUtils.saveToPath(this, path);
    }
}
